from attacker.CPA import CorrcoefCalculator
import numpy as np

def corrcoef(traces,point_num,midvalue):
    if len(midvalue.shape)==2:
        corrcoef_calculator = CorrcoefCalculator()
        corrcoef_calculator.update(traces, midvalue)
        corrcoef = corrcoef_calculator.get_corrcoef()
        corrcoef_point = np.max(np.abs(corrcoef), axis=0)
    else:
        corrcoef_point=np.abs(np.corrcoef(midvalue,traces.transpose())[0,1:])
    corrcoef_sort = np.argsort(-corrcoef_point)
    return corrcoef_sort[:point_num]

def snr(traces, point_num):
    traces_mean = np.mean(traces, axis=0)
    traces_std = np.std(traces, axis=0, ddof=0)
    SNR = np.abs(np.where(traces_std == 0, 0, traces_mean / traces_std))
    SNR_sort = np.argsort(SNR)[::-1]
    return SNR_sort[:point_num]